package com.cyy.trie;

import java.util.HashMap;

/**
 * Trie
 * Trie 也叫做字典树、前缀树（Prefix Tree）、单词查找树
 * Trie 搜索字符串的效率主要跟字符串的长度有关
 * @author 会玩的洋洋
 * @since 2022-03-01
 * @param <V>
 */
public class Trie<V> {
    private int size;

    private Node<V> root;

    /**
     * Trie内元素数量
     * @return int
     */
    public int size() {
        return this.size;
    }

    /**
     * 判断 Trie 内是否为空
     * @return boolean
     */
    public boolean isEmpty() {
        return this.size == 0;
    }

    /**
     * 清空元素
     */
    public void clear() {
        this.size = 0;
        root = null;
    }

    /**
     * 通过key获取value
     * @param key 键
     * @return V
     */
    public V get(String key) {
        Node<V> node = node(key);
        return node != null ? node.value : null;
    }

    /**
     * 判断是否存在 key
     * @param key 键
     * @return boolean
     */
    public boolean contains(String key) {
        Node<V> node = node(key);
        return node != null && node.word;
    }

    /**
     * 添加元素
     * @param key 键
     * @param value 值
     * @return V
     */
    public V add(String key, V value) {
        keyCheck(key);

        if (root == null) {
            root = new Node<>(null);
        }
        Node<V> node = root;
        int len = key.length();
        for (int i = 0; i < len; i++) {
            char c = key.charAt(i);
            boolean emptyChildren = node.children == null;
            Node<V> childNode = emptyChildren ? null : node.children.get(c);
            // 如果子节点不存在 则开辟一条新的子路线
            if (childNode == null) {
                childNode = new Node<>(node);
                childNode.character = c;
                node.children = emptyChildren ? new HashMap<>() : node.children;
                node.children.put(c, childNode);
            }
            node = childNode;
        }
        // 如果找到的node 是一个单词 则覆盖原来的value
        if (node.word) {
            V oldValue = node.value;
            node.value = value;
            return oldValue;
        }

        node.word = true;
        node.value = value;
        ++size;
        return null;
    }

    /**
     * 删除元素
     * @param key 键
     * @return V
     */
    public V remove(String key) {
        Node<V> node = node(key);
        if (node == null || !node.word) {
            return null;
        }
        size--;
        V oldValue = node.value;
        // 如果存在子节点，直接将本节点的word置为false
        if (node.children == null || !node.children.isEmpty()) {
            node.word = false;
            node.value = null;
            return oldValue;
        }

        // 需要从下向上找
        Node<V> parent = null;
        while ((parent = node.parent) != null) {
            parent.children.remove(node.character);
            if (parent.word || !parent.children.isEmpty()) {
                break;
            }
            node = parent;
        }

        return oldValue;
    }

    /**
     * 查看是否存在以 prefix 前缀开头的
     * @param prefix 前缀
     * @return boolean
     */
    public boolean startsWith(String prefix) {
        return node(prefix) != null;
    }

    /**
     * 根据key 查找节点
     * @param key 键
     * @return Node<V>
     */
    private Node<V> node(String key) {
        keyCheck(key);
        Node<V> node = root;
        int len = key.length();
        for (int i = 0; i < len; i++) {
            if (node == null || node.children == null || node.children.isEmpty()) {
                return null;
            }
            char c = key.charAt(i);
            node = node.children.get(c);
        }
        return node;
    }

    private void keyCheck(String key) {
        if (key == null || key.length() == 0) {
            throw new IllegalArgumentException("key not be empty!!");
        }
    }

    private static class Node<V> {
        Node<V> parent;
        HashMap<Character, Node<V>> children;
        Character character;
        V value;
        //是否为单词的结尾（是否为一个完整的单词）
        boolean word;

        public Node(Node<V> parent) {
            this.parent = parent;
        }
    }
}
